#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
文件处理器模块
负责处理不同格式的文档并提取文本内容
"""

import os
import logging
import chardet
from typing import Optional

# 导入文件格式处理库
try:
    import docx  # python-docx for .docx files
    import PyPDF2  # PyPDF2 for .pdf files
    import markdown  # markdown for .md files
    from bs4 import BeautifulSoup  # beautifulsoup4 for .html files
    EXTRA_FORMATS_AVAILABLE = True
except ImportError:
    EXTRA_FORMATS_AVAILABLE = False

# 配置日志
logger = logging.getLogger(__name__)

class FileProcessor:
    """文件处理器类"""
    
    def __init__(self, config=None):
        """
        初始化文件处理器
        
        Args:
            config (dict): 配置参数
        """
        self.config = config or {}
        self.supported_formats = self.config.get('支持格式', ['txt', 'pdf', 'docx', 'md', 'html'])
        self.max_file_size = self.config.get('最大大小', 10 * 1024 * 1024)  # 默认10MB
        
        if not EXTRA_FORMATS_AVAILABLE:
            # 如果没有安装额外的格式支持库，则只支持txt格式
            self.supported_formats = ['txt']
            logger.warning("未安装额外的文件格式支持库，仅支持txt格式")
        
        logger.info(f"文件处理器初始化完成，支持的格式：{self.supported_formats}")
    
    def update_config(self, config):
        """
        更新配置
        
        Args:
            config (dict): 新的配置参数
        """
        self.config = config or {}
        self.supported_formats = self.config.get('支持格式', ['txt', 'pdf', 'docx', 'md', 'html'])
        self.max_file_size = self.config.get('最大大小', 10 * 1024 * 1024)
        
        logger.info("文件处理器配置已更新")
    
    def extract_text(self, file_path: str, file_format: Optional[str] = None) -> str:
        """
        从文件中提取文本
        
        Args:
            file_path (str): 文件路径
            file_format (str, optional): 文件格式，如果未指定则从文件扩展名推断
            
        Returns:
            str: 提取的文本内容
            
        Raises:
            ValueError: 如果文件格式不支持或文件大小超过限制
            FileNotFoundError: 如果文件不存在
            Exception: 其他处理错误
        """
        # 检查文件是否存在
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在：{file_path}")
        
        # 检查文件大小
        file_size = os.path.getsize(file_path)
        if file_size > self.max_file_size:
            raise ValueError(f"文件大小超过限制：{file_size} > {self.max_file_size} bytes")
        
        # 如果未指定格式，从文件扩展名推断
        if file_format is None:
            file_format = os.path.splitext(file_path)[1].lower()[1:]
        
        # 检查格式是否支持
        if file_format not in self.supported_formats:
            raise ValueError(f"不支持的文件格式：{file_format}")
        
        try:
            # 根据文件格式调用相应的处理方法
            if file_format == 'txt':
                return self._process_txt(file_path)
            elif file_format == 'pdf':
                return self._process_pdf(file_path)
            elif file_format == 'docx':
                return self._process_docx(file_path)
            elif file_format == 'md':
                return self._process_markdown(file_path)
            elif file_format == 'html':
                return self._process_html(file_path)
            else:
                raise ValueError(f"未知的文件格式：{file_format}")
        
        except Exception as e:
            logger.error(f"处理文件失败：{str(e)}")
            raise
    
    def _process_txt(self, file_path: str) -> str:
        """处理txt文件"""
        # 首先尝试检测文件编码
        with open(file_path, 'rb') as f:
            raw_data = f.read()
            result = chardet.detect(raw_data)
            encoding = result['encoding']
        
        # 使用检测到的编码读取文件
        try:
            with open(file_path, 'r', encoding=encoding) as f:
                return f.read()
        except UnicodeDecodeError:
            # 如果检测的编码不正确，尝试常用编码
            encodings = ['utf-8', 'gbk', 'gb2312', 'ascii']
            for enc in encodings:
                try:
                    with open(file_path, 'r', encoding=enc) as f:
                        return f.read()
                except UnicodeDecodeError:
                    continue
            
            # 如果所有编码都失败，使用二进制方式读取并解码
            with open(file_path, 'rb') as f:
                content = f.read()
                return content.decode('utf-8', errors='ignore')
    
    def _process_pdf(self, file_path: str) -> str:
        """处理PDF文件"""
        if not EXTRA_FORMATS_AVAILABLE:
            raise ValueError("未安装PDF支持库")
        
        text = []
        with open(file_path, 'rb') as file:
            # 创建PDF文件阅读器
            pdf_reader = PyPDF2.PdfReader(file)
            
            # 遍历所有页面
            for page in pdf_reader.pages:
                try:
                    text.append(page.extract_text())
                except Exception as e:
                    logger.warning(f"提取PDF页面文本失败：{str(e)}")
        
        return "\n".join(text)
    
    def _process_docx(self, file_path: str) -> str:
        """处理Word文档"""
        if not EXTRA_FORMATS_AVAILABLE:
            raise ValueError("未安装Word文档支持库")
        
        doc = docx.Document(file_path)
        text = []
        
        # 提取段落文本
        for paragraph in doc.paragraphs:
            if paragraph.text.strip():
                text.append(paragraph.text)
        
        # 提取表格文本
        for table in doc.tables:
            for row in table.rows:
                row_text = []
                for cell in row.cells:
                    if cell.text.strip():
                        row_text.append(cell.text)
                if row_text:
                    text.append(" | ".join(row_text))
        
        return "\n".join(text)
    
    def _process_markdown(self, file_path: str) -> str:
        """处理Markdown文件"""
        if not EXTRA_FORMATS_AVAILABLE:
            raise ValueError("未安装Markdown支持库")
        
        # 读取markdown文件
        with open(file_path, 'r', encoding='utf-8') as f:
            md_content = f.read()
        
        # 转换为HTML
        html_content = markdown.markdown(md_content)
        
        # 使用BeautifulSoup提取文本
        soup = BeautifulSoup(html_content, 'html.parser')
        
        # 获取所有文本，保留基本格式
        text = []
        for element in soup.descendants:
            if isinstance(element, str) and element.strip():
                text.append(element.strip())
        
        return "\n".join(text)
    
    def _process_html(self, file_path: str) -> str:
        """处理HTML文件"""
        if not EXTRA_FORMATS_AVAILABLE:
            raise ValueError("未安装HTML支持库")
        
        # 读取HTML文件
        with open(file_path, 'r', encoding='utf-8') as f:
            html_content = f.read()
        
        # 使用BeautifulSoup解析HTML
        soup = BeautifulSoup(html_content, 'html.parser')
        
        # 移除script和style元素
        for script in soup(["script", "style"]):
            script.decompose()
        
        # 获取文本
        text = soup.get_text(separator='\n')
        
        # 清理文本
        lines = [line.strip() for line in text.splitlines() if line.strip()]
        
        return "\n".join(lines)
    
    def get_supported_formats(self) -> list:
        """
        获取支持的文件格式列表
        
        Returns:
            list: 支持的文件格式列表
        """
        return self.supported_formats
    
    def check_file_support(self, file_path: str) -> tuple:
        """
        检查文件是否支持处理
        
        Args:
            file_path (str): 文件路径
            
        Returns:
            tuple: (是否支持, 原因)
        """
        # 检查文件是否存在
        if not os.path.exists(file_path):
            return False, "文件不存在"
        
        # 检查文件大小
        file_size = os.path.getsize(file_path)
        if file_size > self.max_file_size:
            return False, f"文件大小超过限制：{file_size} > {self.max_file_size} bytes"
        
        # 检查文件格式
        file_format = os.path.splitext(file_path)[1].lower()[1:]
        if file_format not in self.supported_formats:
            return False, f"不支持的文件格式：{file_format}"
        
        return True, "文件支持处理"